
#include "histogram_tiling.h"
#include "register/op_def_registry.h"

namespace optiling {

constexpr int32_t CORE_NUM = 1;
constexpr int32_t UB_SIZE = 196608;

constexpr int32_t X_INDEX = 0;
constexpr int32_t BINS_ATTR_INDEX = 0;
constexpr int32_t MIN_ATTR_INDEX = 1;
constexpr int32_t MAX_ATTR_INDEX = 2;
constexpr int32_t BYTE_REPEAT = 256;

constexpr int32_t BUF_COUNT_CACHED = (1) * 1 + 2;
constexpr int32_t BUF_COUNT_NOT_CACHED = (1) * 2 + 2;

static ge::graphStatus TilingFunc(gert::TilingContext* context) {
    HistogramTilingData tiling;
    uint32_t totalCoreNum = CORE_NUM;
    uint64_t ubSizePlatForm = UB_SIZE;

    const gert::Shape xShape = context->GetInputShape(X_INDEX)->GetStorageShape();
    const ge::DataType dataType = context->GetInputDesc(X_INDEX)->GetDataType();
    int32_t dTypeSize = 4;
    int32_t perRepeatCount = BYTE_REPEAT / dTypeSize;
    int64_t totalDataCount = xShape.GetShapeSize();
    tiling.set_totalDataCount(totalDataCount);

    uint64_t ubMaxProcCount = ubSizePlatForm / BUF_COUNT_CACHED / BYTE_REPEAT * perRepeatCount;

    context->SetTilingKey(101);
    if (ubMaxProcCount < totalDataCount) {
        printf("==zf==tiling key 201\n");
        context->SetTilingKey(201);
        ubMaxProcCount = ubSizePlatForm / BUF_COUNT_NOT_CACHED / BYTE_REPEAT * perRepeatCount;
        tiling.set_loopTime(totalDataCount / ubMaxProcCount);
        tiling.set_tailCount(totalDataCount % ubMaxProcCount);
    }

    tiling.set_ubMaxProcCount(ubMaxProcCount);
    auto attrs = context->GetAttrs();
    if (attrs == nullptr) {
        return ge::GRAPH_FAILED;
    }
    const int64_t* binsPtr = attrs->GetInt(BINS_ATTR_INDEX);
    tiling.set_binsValue(*binsPtr);
    const float* minPtr = attrs->GetFloat(MIN_ATTR_INDEX);
    float minValue = *minPtr;
    tiling.set_minValue(minValue);
    const float* maxPtr = attrs->GetFloat(MAX_ATTR_INDEX);
    float maxValue = *maxPtr;
    tiling.set_maxValue(maxValue);

    tiling.set_isMinMaxEqual(std::abs(maxValue - minValue) < 1e-5);

    context->SetBlockDim(totalCoreNum);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    printf(
        "==zf==Tiling Data, ubMaxProcCount:%lu, totalDataCount:%lu, binsValue:%ld, minValue:%f, maxValue:%f, "
        "isMinMaxEqual:%ld, loopTime:%lu, tailCount:%lu\n",
        tiling.get_ubMaxProcCount(), tiling.get_totalDataCount(), tiling.get_binsValue(), tiling.get_minValue(),
        tiling.get_maxValue(), tiling.get_isMinMaxEqual(), tiling.get_loopTime(), tiling.get_tailCount());

    // 以下代码是ascendebug必须的
    size_t* workspaces = context->GetWorkspaceSizes(1);
    workspaces[0] = 16 * 1024 * 1024;
    return ge::GRAPH_SUCCESS;
}
}  // namespace optiling

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context) {
    return GRAPH_SUCCESS;
}
}  // namespace ge

namespace ops {
class Histogram : public OpDef {
public:
    explicit Histogram(const char* name) : OpDef(name) {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("bins").AttrType(OPTIONAL).Int(100);
        this->Attr("min").AttrType(OPTIONAL).Float(0.0);
        this->Attr("max").AttrType(OPTIONAL).Float(0.0);

        this->SetInferShape(ge::InferShape);

        this->AICore().SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");
    }
};

OP_ADD(Histogram);
}  // namespace ops
